import sys, struct
import math as pm
import numpy as np
import pylab as p
import math
import csv
import matplotlib.pyplot as plt
import matplotlib as mlt
import h5py
from netcdf_reader import *
from scipy import ndimage
import scipy.sparse.linalg
import nrrd
import matplotlib.patches as mpatches

TOTMEMS = 16
X = 152
Y = 152
DIM = 2

#OUTPUT_DATA_DIR = '/home/behollis/Dropbox/visweek2015/oslines/'
#INPUT_DATA_DIR = '/home/behollis/Dropbox/stirringData/data/'
OUTPUT_DATA_DIR = '/home/brad/visweek2015-revision-code/oslines/'


def showSlines(file):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    
    mag = np.zeros(shape=(X,Y))
    for x in range(0,X):
        for y in range(0,Y):
            avg_vel_mag = 0.0
            for idx in range(7,8):
                avg_vel_mag += math.sqrt( math.pow(vclin[idx][x,y][0],2) + math.pow(vclin[idx][x,y][1],2) ) 
                
            mag[x,y] = avg_vel_mag# / float(TOTMEMS)
    
    plt.imshow(mag, origin='lower', interpolation=None, vmin=0, vmax=4, cmap='gist_ncar')
    plt.colorbar()
    
    '''
    mems = [ d for d in f ]
    for m in mems:
        x = '/x020'
        y = '/y120'
        sl = f[m+x+y]
        plt.plot(sl[0], sl[1], color ='white')
        print len(sl[0])
        
    plt.plot(20., 120., '+', color='red', linewidth=50.0, ms=10.0) # draw seed loc
    '''

    plt.show()
    
    f.close()

def showFTVA(file):
    
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
   
    ftva = np.zeros(shape=(X,Y))
   
    for x in range(0,X):
        for y in range(0, Y):
            c = f['cov'][x,y]
            
            #try:
            w, v = np.linalg.eig(c)
            sorted(w)
            #except:
            #    w[0] = 0
                
            ftva[x,y] = w[0]
    
    plt.imshow(ftva, origin='lower', interpolation=None, vmin=0, vmax=2000)
    plt.colorbar()
    plt.show()
    f.close()

def showEntropySlineClusters(file,mult,x=None,y=None):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['slineclusters'] 
    
    max = int(np.amax(termClusters))
    
    FACT = mult
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*FACT*(max+1),max*FACT)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))

    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    fig = plt.figure()
    
    imgplot = plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)
    #plt.colorbar()
    
    tlabs = [str(i) for i in range(0,max+1)]
    tks = [i + 0.5 for i in range(0,max+1)]
    
    
    #plt.colorbar()
    cb = fig.colorbar(imgplot, ticks=tks)
    cb.ax.set_yticklabels(tlabs)
    
    '''
    patches = list()    
    for idx in range(0,len(cmaplist)):
        patches.append( mpatches.Patch(color='red', label=str(idx)+' clusters') )
        
    plt.legend(handles=patches)
    '''
    
    '''
    #plot slines at seed
    if x is not None:
        f = h5py.File(OUTPUT_DATA_DIR + 'oceanSlinesLevel00.hdf5', 'r')
        mems = [ d for d in f ]
        
        for m in mems:
            dir = m + '/x' + str(y).zfill(3) + '/y' + str(x).zfill(3)
            sl = f[dir]
            plt.plot( sl[1], sl[0] )
            
        plt.plot(x, y, '+', color='white', linewidth=10.0, ms=10.0) # draw seed loc
    '''
    
    plt.show()
    f.close()

def showTerminalParticleClusters(file, mult=10):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['ptclusters'] 
    
    max = int(np.amax(termClusters))
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*mult*(max+1),max*mult)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))
    
    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    imgplot = plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)
    
    '''
    ps = list()
    
    for idx in range( 0,len(cmaplist) ):
        p, = mpatches.Patch(color='red', label='2' )
        ps.append(p) 
        
    plt.legend(handles=ps)
    '''
    
    plt.colorbar()
    
    plt.show()
    f.close()
    
def showEntropyFields(file):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    
    avgAngEntr = f['avgAngularEntropy']
    avgAngEntrMin = np.min(avgAngEntr)
    avgAngEntrMax = np.max(avgAngEntr)
    avgLinEntr = f['avgLinearEntropy']
    avgLinEntrMin = np.min(avgLinEntr)
    avgLinEntrMax = np.max(avgLinEntr)
    
    
    '''
    dsampled = output.create_dataset(name='ptsSampledForClustering', shape=(LAT,LON), dtype='f')
    dentroL = output.create_dataset(name='avgAngularEntropy', shape=(LAT,LON), dtype='f')
    dentroA = output.create_dataset(name='avgLinearEntropy', shape=(LAT,LON), dtype='f')
    '''

    cmap = plt.cm.jet
    
#    plt.imshow(avgAngEntr, origin='lower', interpolation=None, vmin=0, vmax=0.05)
    plt.imshow(avgAngEntr, origin='lower', interpolation=None, vmin=0.0, vmax=1.0)
    cb = plt.colorbar(); 
    plt.show()
    
#    plt.imshow(avgLinEntr, origin='lower', interpolation=None, vmin=0, vmax=0.5)
    plt.imshow(avgLinEntr, origin='lower', interpolation=None, vmin=0.0, vmax=1.0)
    plt.colorbar(); 
    plt.show()
    
    powOf2 = np.zeros(shape=avgLinEntr.shape); powOf2.fill(2.0)
    sqRoot = np.zeros(shape=avgLinEntr.shape); sqRoot.fill(0.5)
    
    gradLinEntr = np.gradient(avgLinEntr)
    gradLinEntrMag = np.power( np.power(gradLinEntr[0], powOf2) + np.power(gradLinEntr[1], powOf2), sqRoot)
    
    gradAngEntr = np.gradient(avgAngEntr)
    gradAngEntrMag = np.power( np.power(gradAngEntr[0], powOf2) + np.power(gradAngEntr[1], powOf2), sqRoot)
    
    # GRADIENT of LINEAR ENTROPY map
    plt.imshow(gradAngEntrMag, origin='lower', interpolation=None, vmin=np.min(gradAngEntrMag), vmax=0.15, cmap=cmap)
    plt.colorbar();
    plt.show()
    
    # GRADIENT of A ENTROPY map
    plt.imshow(gradLinEntrMag, origin='lower', interpolation=None, vmin=np.min(gradLinEntrMag), vmax=0.15, cmap=cmap)
    plt.colorbar(); 
    plt.show()
    
    
    
    f.close()
    
def showSlineSamplingRate(file, mult=10):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['ptsSampledForClustering'] 
    
    max = int(np.amax(termClusters))
    min = int(np.amax(termClusters))
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*mult*(max+1),max*mult)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))
    
    fig, ax = plt.subplots()
    
    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    bounds = bounds[3:-1]
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    imgplot = plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)#, vmin = 3, vmax = max)
    
    '''
    ps = list()
    
    for idx in range( 0,len(cmaplist) ):
        p, = mpatches.Patch(color='red', label='2' )
        ps.append(p) 
        
    plt.legend(handles=ps)
    '''
    tlabs = [str(i) for i in range(3,max+1)]
    tks = [i + 0.5 for i in range(3,max+1)]
    
    
    #plt.colorbar()
    cb = fig.colorbar(imgplot, ticks=tks)
    cb.ax.set_yticklabels(tlabs)
    
    plt.show()
    f.close()
    
if __name__ == '__main__':
    
    '''
    vclin = np.zeros( shape = (TOTMEMS, X, Y, DIM) )
    for idx in range(0,TOTMEMS):
        dir = 'run'+str(idx).zfill(1) +'/vel050.nrrd'
        
        vfile = nrrd.read(INPUT_DATA_DIR + dir)[0][:, 0:X, 0:Y]
        vfile = np.swapaxes(vfile, 0, 1)
        vfile = np.swapaxes(vfile, 1, 2)
        vclin[idx][:] = vfile[:]
    '''
    
    
#    showSlines('industrialStirringSlinesTS050.500steps.ss05.2.hdf5')
#    showFTVA(file='stirringFtvaForward.hdf5')
#    showFTVA(file='stirringFtvaBackward.hdf5') 
#    showTerminalParticleClusters(file='stirTermClusFor.hdf5', mult=7)
#    showTerminalParticleClusters(file='stirTermClusBack.hdf5', mult=5)
    showSlineSamplingRate(file='stir.SlineClusters.varSampling.3to13pts.hdf5',mult=2)
    showEntropySlineClusters(file='stir.SlineClusters.varSampling.3to13pts.hdf5',mult=8)
#    showEntropyFields(file='stirringEntropiesONLY.fixed.hdf5')
    
    
    
    
    
    
    